package etxt2db.api;
import java.io.BufferedWriter;
import java.io.File;
import java.io.FileNotFoundException;
import java.io.FileWriter;
import java.io.IOException;
import java.io.Writer;
import java.text.ParseException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;

import edu.cmu.minorthird.text.MutableTextLabels;
import edu.cmu.minorthird.text.RegexTokenizer;
import edu.cmu.minorthird.text.Span;
import etxt2db.corpus.CompleteTextBaseLoader;
import etxt2db.corpus.MyTextBaseLoader;

/**
 * <p>
 * This class represents an object that contains functions to execute an Information Extraction
 * classifier with a given text.
 * </p>
 *
 * @author Gon�alo Sim�es
 */

public class ClassificationExecutor {
	
	private boolean docPerFile = true;
	private String beginTagString = "<";
	private String endTagString = ">";
	
	public ClassificationExecutor(){
	}
	
	
	/**
     * Function that prints the spans that were classified with a given set of attributes
     * @param file A File containing the text you want to classify
     * @param annotator The classifier you will use
     * @param attributes The list containing the attributes you want to show
     */
	public void printClassificationResults(File file, ClassificationModel annotator, List<String> attributes) 
		throws IOException, ParseException{
		MutableTextLabels labels2 = annotateTestingDocument(file, annotator);
		if(attributes == null){
			attributes = new ArrayList<String>(labels2.getTypes());
		}
		printAnnotationResults(labels2, attributes);
	}
	
	/**
     * Function that returns the classified spans from a text present in a String 
     * as a Map in which the key is a class and the value is a list of segments from 
     * the text.
     * @param str A String containing the text you want to classify
     * @param annotator The classifier you will use
     * @param attributes The list containing the attributes you want to show
     * @return Map in which the key is a class and the value is a list of segments from the text
     */
	public Map<String,List<String>> getClassifiedSegments(String str, ClassificationModel annotator, List<String> attributes) 
	throws IOException, ParseException{
		File outputFile = File.createTempFile("String", ".tmp");
		Writer output = new BufferedWriter(new FileWriter(outputFile));
		try {
			//FileWriter always assumes default encoding is OK!
			output.write( str );
		}
		finally {
			output.close();
		}
		
		CompleteTextBaseLoader baseLoader2 = new CompleteTextBaseLoader();
		if(!docPerFile){
			baseLoader2.setDocumentStyle(MyTextBaseLoader.DOC_PER_LINE);
		}
		baseLoader2.setBeginTag(beginTagString);
		baseLoader2.setEndTag(endTagString);
		baseLoader2.load(outputFile,new RegexTokenizer("\\s*([0-9]+|[a-zA-Z��������������������]+|\\W)\\s*"));
		MutableTextLabels labels2 = baseLoader2.getLabels();
		//BasicTextBase base2 = new BasicTextBase(new RegexTokenizer("\\s*([0-9]+|[a-zA-Z��������������������]+|\\W)\\s*"));
		//base2.loadDocument("justString", str);
		//MutableTextLabels labels2 = new BasicTextLabels(base2);
		annotator.annotate(labels2);
		if(attributes == null){
			attributes = new ArrayList<String>(labels2.getTypes());
		}
		return classifySegments(labels2,attributes);
	}
	
	public Map<String,List<ClassifiedSpan>> getClassifiedSpans(String str, ClassificationModel annotator, List<String> attributes) 
	throws IOException, ParseException{
		File outputFile = File.createTempFile("String", ".tmp");
		Writer output = new BufferedWriter(new FileWriter(outputFile));
		try {
			//FileWriter always assumes default encoding is OK!
			output.write( str );
		}
		finally {
			output.close();
		}
		
		CompleteTextBaseLoader baseLoader2 = new CompleteTextBaseLoader();
		if(!docPerFile){
			baseLoader2.setDocumentStyle(MyTextBaseLoader.DOC_PER_LINE);
		}
		baseLoader2.setBeginTag(beginTagString);
		baseLoader2.setEndTag(endTagString);
		baseLoader2.load(outputFile,new RegexTokenizer("\\s*([0-9]+|[a-zA-Z��������������������]+|\\W)\\s*"));
		MutableTextLabels labels2 = baseLoader2.getLabels();
		//BasicTextBase base2 = new BasicTextBase(new RegexTokenizer("\\s*([0-9]+|[a-zA-Z��������������������]+|\\W)\\s*"));
		//base2.loadDocument("justString", str);
		//MutableTextLabels labels2 = new BasicTextLabels(base2);
		annotator.annotate(labels2);
		if(attributes == null){
			attributes = new ArrayList<String>(labels2.getTypes());
		}
		return classifySpans(labels2,attributes);
	}
	
	/**
     * Function that returns the classified spans from a text present in a String 
     * as a Map in which the key is a class and the value is a list of segments from 
     * the text.
     * @param str A String containing the text you want to classify
     * @param annotator The classifier you will use
     * @param attributes The list containing the attributes you want to show
     * @return Map in which the key is a class and the value is a list of segments from the text
     */
	public String getClassifiedString(String str, ClassificationModel annotator, List<String> attributes) 
	throws IOException, ParseException{
		File outputFile = File.createTempFile("String", ".tmp");
		Writer output = new BufferedWriter(new FileWriter(outputFile));
		try {
			//FileWriter always assumes default encoding is OK!
			output.write( str );
		}
		finally {
			output.close();
		}
		
		CompleteTextBaseLoader baseLoader2 = new CompleteTextBaseLoader();
		if(!docPerFile){
			baseLoader2.setDocumentStyle(MyTextBaseLoader.DOC_PER_LINE);
		}
		baseLoader2.setBeginTag(beginTagString);
		baseLoader2.setEndTag(endTagString);
		baseLoader2.load(outputFile,new RegexTokenizer("\\s*([0-9]+|[a-zA-Z��������������������]+|\\W)\\s*"));
		MutableTextLabels labels2 = baseLoader2.getLabels();
		annotator.annotate(labels2);
		
		//BasicTextBase base2 = new BasicTextBase(new RegexTokenizer("\\s*([0-9]+|[a-zA-Z��������������������]+|\\W)\\s*"));
		//base2.loadDocument("justString", str);
		//MutableTextLabels labels2 = new BasicTextLabels(base2);
		if(attributes == null){
			attributes = new ArrayList<String>(labels2.getTypes());
		}
		return getAnnotationResultsInString(labels2, attributes);
	}
	
	/**
     * Function that returns the classified spans from a text present in a File
     * as a Map in which the key is a class and the value is a list of segments from 
     * the text.
     * @param file A File containing the text you want to classify
     * @param annotator The classifier you will use
     * @param attributes The list containing the attributes you want to show
     * @return Map in which the key is a class and the value is a list of segments from the text
     */
	public Map<String,List<String>> getClassifiedSegments(File file, ClassificationModel annotator, List<String> attributes) 
	throws IOException, ParseException{
		MutableTextLabels labels2 = annotateTestingDocument(file, annotator);
		if(attributes == null){
			attributes = new ArrayList<String>(labels2.getTypes());
		}
		return classifySegments(labels2,attributes);
	}
	
	public Map<String,List<ClassifiedSpan>> getClassifiedSpans(File file, ClassificationModel annotator, List<String> attributes) 
	throws IOException, ParseException{
		MutableTextLabels labels2 = annotateTestingDocument(file, annotator);
		if(attributes == null){
			attributes = new ArrayList<String>(labels2.getTypes());
		}
		return classifySpans(labels2,attributes);
	}
	
	/**
     * Function that generates an annotated file
     * @param file A File containing the text you want to classify
     * @param annotator The classifier you will use
     * @param attributes The list containing the attributes you want to show
     * @param outputDirectory The output file
     */
	public void createClassifiedFile(File file, ClassificationModel annotator, List<String> attributes, String outputDirectory) 
		throws IOException, ParseException{
		MutableTextLabels labels2 = annotateTestingDocument(file, annotator);
		if(attributes == null){
			attributes = new ArrayList<String>(labels2.getTypes());
		}
		writeAnnotationResultsToFile(labels2,attributes,outputDirectory);
	}
	
	/**
     * Function that generates an annotated string
     * @param file A File containing the text you want to classify
     * @param annotator The classifier you will use
     * @param attributes The list containing the attributes you want to show
     * @return String containing classifiedContents
     */
	public String getClassifiedString(File file, ClassificationModel annotator, List<String> attributes) 
		throws IOException, ParseException{
		MutableTextLabels labels2 = annotateTestingDocument(file, annotator);
		if(attributes == null){
			attributes = new ArrayList<String>(labels2.getTypes());
		}
		return getAnnotationResultsInString(labels2,attributes);
	}
	
	private String getAnnotationResultsInString(MutableTextLabels labels2, List<String> attributes) throws IOException{
		String result = "";
		Iterator<Span> iter = labels2.getTextBase().documentSpanIterator();
		while(iter.hasNext()){
			Span span = iter.next();
			result = result + "-----" +  span.getDocumentId() + "-----" + "\n";
			String document = span.asString();
			List<ListEntry> tagsList = new ArrayList<ListEntry>();
			for(String attribute : attributes){
				for(Span text : labels2.getTypeSet(attribute, span.getDocumentId())){
					int loChar = text.getLoChar()-span.getLoChar();
    				int hiChar = text.getHiChar()-span.getLoChar();
					tagsList.add(new ListEntry(attribute,true,loChar,hiChar));
					tagsList.add(new ListEntry(attribute,false,hiChar,loChar));
				}
			}
			Collections.sort(tagsList);
			for(ListEntry le : tagsList){
				String tag;
				if(le.isOpening){
					tag = this.beginTagString + le.tagName + this.endTagString;
				}else{
					tag = this.beginTagString + "/" + le.tagName + this.endTagString;
				}
				document = document.substring(0,le.index) + tag + document.substring(le.index);
			}
			result = result + document + "\n";
		}
		
		return result;
	}
	
	private void writeAnnotationResultsToFile(MutableTextLabels labels2, List<String> attributes, String outputDirectory) throws IOException{
		Iterator<Span> iter = labels2.getTextBase().documentSpanIterator();
		while(iter.hasNext()){
			Span span = iter.next();
			String document = span.asString();
			List<ListEntry> tagsList = new ArrayList<ListEntry>();
			for(String attribute : attributes){
				for(Span text : labels2.getTypeSet(attribute, span.getDocumentId())){
					int loChar = text.getLoChar()-span.getLoChar();
    				int hiChar = text.getHiChar()-span.getLoChar();
					tagsList.add(new ListEntry(attribute,true,loChar,hiChar));
					tagsList.add(new ListEntry(attribute,false,hiChar,loChar));
				}
			}
			Collections.sort(tagsList);
			for(ListEntry le : tagsList){
				String tag;
				if(le.isOpening){
					tag = this.beginTagString + le.tagName + this.endTagString;
				}else{
					tag = this.beginTagString + "/" + le.tagName + this.endTagString;
				}
				document = document.substring(0,le.index) + tag + document.substring(le.index);
			}
			String outputDir = outputDirectory;
			File outputFile = new File(outputDir, span.getDocumentId());
			outputFile.createNewFile();
			writeTaggedFile(outputFile, document);
		}
	}
	
	private void printAnnotationResults(MutableTextLabels labels2, List<String> attributes){
		for(String attribute : attributes){
			System.out.println("The texts classified as \"" + attribute + "\" are:");
			Iterator<Span> iter = labels2.getTextBase().documentSpanIterator();
			while(iter.hasNext()){
				Span span = iter.next();
				for(Span text : labels2.getTypeSet(attribute, span.getDocumentId())){
					System.out.println(text.asString());
				}
			}
		}
	}
	
	private Map<String,List<String>> classifySegments(MutableTextLabels labels2, List<String> attributes){
		Map<String,List<String>> strings = new HashMap<String,List<String>>();
		for(String attribute : attributes){
			List<String> list = new ArrayList<String>();
			Iterator<Span> iter = labels2.getTextBase().documentSpanIterator();
			while(iter.hasNext()){
				Span span = iter.next();
				for(Span text : labels2.getTypeSet(attribute, span.getDocumentId())){
					list.add(text.asString());
				}
			}
			strings.put(attribute, list);
		}
		return strings;
	}
	
	private Map<String,List<ClassifiedSpan>> classifySpans(MutableTextLabels labels2, List<String> attributes){
		Map<String,List<ClassifiedSpan>> strings = new HashMap<String,List<ClassifiedSpan>>();
		Iterator<Span> iter = labels2.getTextBase().documentSpanIterator();
		while(iter.hasNext()){
			Span span = iter.next();
			String document = span.asString();
			List<ClassifiedSpan> tagsList = strings.get(span.getDocumentId());
			if(tagsList==null){
				tagsList = new ArrayList<ClassifiedSpan>();
			}
			for(String attribute : attributes){
				for(Span text : labels2.getTypeSet(attribute, span.getDocumentId())){
					int loChar = text.getLoChar();
    				int hiChar = text.getHiChar();
					ClassifiedSpan s = new ClassifiedSpan(loChar,hiChar,attribute);
					tagsList.add(s);
				}
			}
			strings.put(span.getDocumentId(), tagsList);
		}
		return strings;
	}
	
	/*private MutableTextLabels annotateString(String str, Annotator annotator){
		MutableTextLabels labels2 = new BasicTextLabels(str);
		System.out.println(labels2.getTextBase());
		annotator.annotate(labels2);
		return labels2;
	}*/
	
	private MutableTextLabels annotateTestingDocument(File file, ClassificationModel annotator) throws IOException, ParseException{
		CompleteTextBaseLoader baseLoader2 = new CompleteTextBaseLoader();
		if(!docPerFile){
			baseLoader2.setDocumentStyle(MyTextBaseLoader.DOC_PER_LINE);
		}
		baseLoader2.setBeginTag(beginTagString);
		baseLoader2.setEndTag(endTagString);
		baseLoader2.load(file,new RegexTokenizer("\\s*([0-9]+|[a-zA-Z��������������������]+|\\W)\\s*"));
		MutableTextLabels labels2 = baseLoader2.getLabels();
		annotator.annotate(labels2);
		return labels2;
	}
	
	private void writeTaggedFile(File aFile, String aContents)
		throws FileNotFoundException, IOException {
		if (aFile == null) {
			throw new IllegalArgumentException("File should not be null.");
		}
		if (!aFile.canWrite()) {
			throw new IllegalArgumentException("File cannot be written: " + aFile);
		}
		
		//use buffering
		Writer output = new BufferedWriter(new FileWriter(aFile));
		try {
			//FileWriter always assumes default encoding is OK!
			output.write( aContents );
		}
		finally {
			output.close();
		}
	}
	
	private class ListEntry implements Comparable<ListEntry>{
		public ListEntry(String tagName, boolean isOpening, int index, int pairIndex){
			this.tagName = tagName;
			this.isOpening = isOpening;
			this.index = index;
			this.pairIndex = pairIndex;
		}
		
		String tagName;
		boolean isOpening;
		int index;
		int pairIndex;
		
		public int compareTo(ListEntry o) {
			if(index>o.index){
				return -1;
			}
			if(index<o.index){
				return 1;
			}
			if(isOpening == true && o.isOpening == false){
				return -1;
			}
			if(isOpening == false && o.isOpening == true){
				return 1;
			}
			if(isOpening == true && pairIndex<o.pairIndex){
				return -1;
			}
			if(isOpening == true && pairIndex>o.pairIndex){
				return 1;
			}
			if(isOpening == false && pairIndex<o.pairIndex){
				return 1;
			}
			if(isOpening == false && pairIndex>o.pairIndex){
				return -1;
			}
			if(isOpening == true){
				return tagName.compareTo(o.tagName);
			}else{
				return -tagName.compareTo(o.tagName);
			}
		}
		
		public String toString(){
			return "[" + tagName + " " + isOpening + " " + index + " " + pairIndex + "]";
		}
	}
	
	public void setBeginTag(String str){
		this.beginTagString = str;
	}
	
	public void setEndTag(String str){
		this.endTagString = str;
	}
	
	public void setDocPerFile(boolean docPerFile){
		this.docPerFile = docPerFile;
	}
	
	public class ClassifiedSpan{
		private String type;
		private int start;
		private int end;
		
		public ClassifiedSpan(int start, int end, String type) {
			super();
			this.start = start;
			this.end = end;
			this.type = type;
		}
		
		public int getStart() {
			return start;
		}
		public void setStart(int start) {
			this.start = start;
		}
		public int getEnd() {
			return end;
		}
		public void setEnd(int end) {
			this.end = end;
		}
		public String getType() {
			return type;
		}
		public void setType(String type) {
			this.type = type;
		}
		
		public String toString(){
			return type + "(" + start + "," + end + ")";
		}
	}
}
